import torch
import torchvision.transforms
from PIL import Image
from model import resnet
import logging
from datetime import datetime

import cv2
import os


def prediction():
    os.chdir("./")  # 日志写入地址
    logging.basicConfig(filename=f'example_{datetime.now().microsecond}.log', level=logging.INFO,
                        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    normalize = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transform = torchvision.transforms.Compose([
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.RandomCrop(224, 4),
        torchvision.transforms.ToTensor(),
        normalize,
    ])
    correct_num = 0
    all_num = 0
    img_dir = "./prediction"
    class_list = os.listdir(img_dir)
    for class_id in class_list:
        class_dir = os.path.join(img_dir, class_id)
        name_list = os.listdir(class_dir)
        all_num = len(os.listdir(class_dir))
        logging.info(f'路径{class_dir},数量{all_num}')
        print(f'路径{class_dir},数量{all_num}')
        for name_id in name_list:
            name_dir = os.path.join(class_dir, name_id)
            # name_dir = './data/section_classification/1/P1-1_image_3.png'
            img = cv2.imread(name_dir, 0)
            img_PIL = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            image = transform(img_PIL)
            image = torch.unsqueeze(image, dim=0)
            # image = torch.reshape(image, (1, 3, 224, 224))
            module = torch.load("./model.pth")
            model = torch.nn.DataParallel(resnet.__dict__['resnet32']())
            model.load_state_dict(module['state_dict'])
            model.cuda()
            model.eval()
            with torch.no_grad():
                output = model(image).cpu()
                predict = torch.softmax(output, dim=1)
                predict_cla = torch.argmax(predict).numpy()
                print(f'切面{name_dir},类型{class_id},预测{predict_cla + 1}')
                if predict_cla + 1 == int(class_id):
                    correct_num += 1

        logging.info(f"类别{class_id}的准确率是{correct_num / all_num}")
        correct_num = 0


def test():
    pass


prediction()
